#!/usr/bin/env python3
import os
import argparse
from pathlib import Path


def list_directory_contents(directory_path, output_file):
    """
    递归遍历目录，将文件路径按目录结构组织并写入txt文件
    """
    directory = Path(directory_path)
    
    if not directory.exists():
        print(f"错误：目录 '{directory_path}' 不存在")
        return
    
    if not directory.is_dir():
        print(f"错误：'{directory_path}' 不是目录")
        return
    
    # 获取目录的绝对路径
    abs_path = directory.resolve()
    
    with open(output_file, 'w', encoding='utf-8') as f:
        # 写入根目录信息
        f.write(f"目录：{abs_path}\n")
        f.write("=" * 80 + "\n\n")
        
        # 递归遍历目录
        for root, dirs, files in os.walk(abs_path):
            # 计算相对路径用于显示
            rel_path = Path(root).relative_to(abs_path)
            
            # 写入当前目录路径作为分割线
            if rel_path == Path('.'):
                f.write(f"【根目录】\n")
            else:
                f.write(f"【{rel_path}】\n")
            f.write("-" * 60 + "\n")
            
            # 写入子目录
            if dirs:
                f.write("子目录：\n")
                for dir_name in sorted(dirs):
                    f.write(f"  📁 {dir_name}\n")
                f.write("\n")
            
            # 写入文件
            if files:
                f.write("文件：\n")
                for file_name in sorted(files):
                    file_path = Path(root) / file_name
                    f.write(f"  📄 {file_name} ({file_path})\n")
                f.write("\n")
            
            f.write("\n")
    
    print(f"目录结构已保存到：{output_file}")


def main():
    # 设置命令行参数
    parser = argparse.ArgumentParser(description='将目录结构写入txt文件')
    parser.add_argument('directory', help='要遍历的目录路径')
    parser.add_argument('-o', '--output', 
                       default='directory_listing.txt',
                       help='输出文件名 (默认: directory_listing.txt)')
    
    args = parser.parse_args()
    
    # 执行目录遍历
    list_directory_contents(args.directory, args.output)


if __name__ == "__main__":
    main()